#define _FILE_OFFSET_BITS 64

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <getopt.h>
#include <stdarg.h>
#include <fcntl.h>
#include <sys/ioctl.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <errno.h>
#include <linux/unistd.h>

#include "loop.h"
#include "bitmap.c"

#define LOOPMAJOR 7

int
is_loop_device (const char *device) {
	struct stat statbuf;

	return (stat(device, &statbuf) == 0 &&
		S_ISBLK(statbuf.st_mode) &&
		major(statbuf.st_rdev) == LOOPMAJOR);
}


int verbose = 0;
char *progname;
char *device;
char *mapfile;
char *backfile;
int block_index = -1;

static void usage(void)
{
	fprintf(stderr, "usage:\n\
  %s loop_device                    # give info\n\
  %s -w loop_device                 # start write limit on device\n\
  %s -x loop_device <mapfile>       # stop and save write limit map to file\n\
  %s -r loop_device <mapfile>       # start read limit\n\
  %s -s loop_device                 # stop read limit on device\n\
  %s -t loop_device -b <n> <file>   # sync block <n> of <file> to device\n\
  %s -c loop_device <file>          # switch back file\n\
  %s -d procfs_rblocked             # show rblocked index\n\
  %s -m wlimit.map                  # show map file info\n",
		progname, progname, progname, progname, progname, progname, progname, progname, progname);
	exit(1);
}

enum {
	CMD_INFO,
	CMD_WLIMIT_START,
	CMD_WLIMIT_STOP,
	CMD_RLIMIT_START,
	CMD_RLIMIT_STOP,
	CMD_SYNC_BLOCK,
	CMD_MAP_INFO,
	CMD_CHANGE,
	CMD_SHOW_RBLOCKED,
};

void show_info()
{
	if (!is_loop_device(device)) {
		fprintf(stderr, "%s is not a loop device\n", device);
		exit(1);
	}

	printf("no information\n");
}

void show_map_info()
{
	struct bitmap_save save;
	struct bitmap* bm;
	int off, fd, ret;
	u32 pos;

	if ((fd = open(mapfile, O_RDONLY)) < 0) {
		perror("open map");
		exit(1);
	}

	read(fd, &save.bitmap_size, 4);
	read(fd, &save.buf_size, 4);
	printf("bitmap size = %d\n", save.bitmap_size);
	printf("encoding buf size = %d\n", save.buf_size);

	if (save.buf_size > 0) {
		save.buf = malloc(save.buf_size);
		if (!save.buf) {
			perror("malloc buf");
			exit(1);
		}
	}

	off = 0;
	while (off < save.buf_size) {
		ret = read(fd, save.buf + off, save.buf_size - off);
		if (ret < 0) {
			perror("read mapfile");
			exit(1);
		}
		off += ret;
			
	}

	bm = bm_alloc(save.bitmap_size);
	if (!bm) {
		perror("alloc struct bitmap");
		exit(1);
	}

	ret = bm_dec_buf(bm, (u32*)save.buf, save.buf_size / 4);
	if (ret != save.buf_size / 4) {
		fprintf(stderr, "error decoding bitmap, ret = %d\n", ret);
		exit(1);
	}
	
	printf("count = %d\n", bm_count(bm));
	if (!verbose)
		goto end;

	pos = 0;
	while(bm_nextb(bm, &pos)) {
		printf("block no. %d\n", pos);
		++pos;
	}

end:
	close(fd);
}

void wlimit_start()
{
	int ret, fd;

	if (!is_loop_device(device)) {
		fprintf(stderr, "%s is not a loop device\n", device);
		exit(1);
	}

	if ((fd = open(device, O_RDONLY)) < 0) {
		perror("open loop");
		exit(1);
	}

	ret = ioctl(fd, LOOP_WLIMIT_START, 0);
	if (ret) {
		perror("set write limit on loop");
		exit(1);
	}
	printf("set write limit: success\n");
}

void wlimit_stop()
{
	int ret, fd, mapfd;
	struct bitmap_save save;
	int off;

	if (!is_loop_device(device)) {
		fprintf(stderr, "%s is not a loop device\n", device);
		exit(1);
	}

	if ((fd = open(device, O_RDONLY)) < 0) {
		perror("open loop");
		exit(1);
	}

	if ((mapfd = open(mapfile, O_WRONLY | O_CREAT | O_TRUNC, S_IRWXU)) < 0) {
		perror("open mapfile");
		exit(1);
	}

	memset(&save, 0, sizeof(save));
	ret = ioctl(fd, LOOP_WLIMIT_STOP, &save);
	if (ret == -1 && errno != EAGAIN) {
		perror("stop write limit phase 1 failed");
		exit(1);
	}
	
	if (save.buf_size > 0) {
		save.buf = malloc(save.buf_size);
		if (!save.buf) {
			printf("buf_size = %d, bitmap_size = %d", save.buf_size, save.bitmap_size);
			perror("malloc buf");
			exit(1);
		}
	}

	ret = ioctl(fd, LOOP_WLIMIT_STOP, &save);
	if (ret) {
		perror("stop write limit phase 2 failed");
		exit(1);
	}

	write(mapfd, &save.bitmap_size, 4);
	write(mapfd, &save.buf_size, 4);

	off = 0;
	while (off < save.buf_size) {
		ret = write(mapfd, save.buf + off, save.buf_size - off);
		if (ret <= 0) {
			perror("write mapfile");
			exit(1);
		}
		off += ret;
	}
	close(fd);
	close(mapfd);
	printf("save write map: success\n");
}

void rlimit_stop()
{
	int ret, fd;

	if (!is_loop_device(device)) {
		fprintf(stderr, "%s is not a loop device\n", device);
		exit(1);
	}

	if ((fd = open(device, O_RDONLY)) < 0) {
		perror("open loop");
		exit(1);
	}

	ret = ioctl(fd, LOOP_RLIMIT_STOP, 0);
	if (ret) {
		perror("stop read limit on loop");
		exit(1);
	}
	printf("stop read limit: success\n");
}

void rlimit_start()
{
	int ret, fd, mapfd;
	struct bitmap_save save;
	int off;

	if (!is_loop_device(device)) {
		fprintf(stderr, "%s is not a loop device\n", device);
		exit(1);
	}

	if ((fd = open(device, O_RDONLY)) < 0) {
		perror("open loop");
		exit(1);
	}

	if ((mapfd = open(mapfile, O_RDONLY)) < 0) {
		perror("open mapfile");
		exit(1);
	}

	read(mapfd, &save.bitmap_size, 4);
	read(mapfd, &save.buf_size, 4);

	if (save.buf_size > 0) {
		save.buf = malloc(save.buf_size);
		if (!save.buf) {
			perror("malloc buf");
			exit(1);
		}
	}

	off = 0;
	while (off < save.buf_size) {
		ret = read(mapfd, save.buf + off, save.buf_size - off);
		if (ret <= 0) {
			perror("read mapfile");
			exit(1);
		}
		off += ret;
	}

	struct bitmap* bm = bm_alloc(save.bitmap_size);
	ret = bm_dec_buf(bm, (u32*)save.buf, save.buf_size / 4);
	if (ret != save.buf_size / 4) {
		fprintf(stderr, "error decoding bitmap, ret = %d\n", ret);
		exit(1);
	}
	
	printf("count = %d\n", bm_count(bm));

	ret = ioctl(fd, LOOP_RLIMIT_START, &save);
	if (ret) {
		perror("error start read limit on loop");
		exit(1);
	}

	close(fd);
	close(mapfd);
	printf("start read limit: success\n");
}

void sync_block()
{
	int ret, fd, backfd;
	int off;
	off_t pos;
	struct loop_block lb;
	u8 block[LOOP_BLOCK_SIZE];
	int isloop = 0;
	int fdmod = O_WRONLY;


	if (is_loop_device(device)) {
		printf("%s is a loop device, use ioctl\n", device);
		isloop = 1;
		fdmod = O_RDONLY;
	}

	if ((fd = open(device, fdmod)) < 0) {
		perror("open loop");
		exit(1);
	}

	if ((backfd = open(backfile, O_RDONLY)) < 0) {
		perror("open backfile");
		exit(1);
	}
	lb.index = block_index;
	lb.len = LOOP_BLOCK_SIZE;
	lb.buf = block;

	pos = ((off_t) lb.index) * LOOP_BLOCK_SIZE;
	ret = lseek(backfd, pos, SEEK_SET);
	if (ret < 0) { 
		perror("lseek");
		exit(1);
	}

	off = 0;
	while (off < lb.len) {
		ret = read(backfd, lb.buf + off, lb.len - off);
		if (ret <= 0) {
			printf("off = %d\n", off);
			perror("read block");
			exit(1);
		}
		off += ret;
	}

	if (isloop)
		ret = ioctl(fd, LOOP_SYNC_BLOCK, &lb);
	else
		ret = write(fd, &lb, sizeof(&lb));

	if (ret) {
		perror("sync block failed");
		exit(1);
	}
	
	close(backfd);
	close(fd);
	printf("sync block: success\n");
}

void change_fd()
{
	int ret, fd, backfd;
	struct loop_info64 lo_info;

	if (!is_loop_device(device)) {
		fprintf(stderr, "%s is not a loop device!\n", device);
		exit(1);
	}

	if ((fd = open(device, O_RDONLY)) < 0) {
		perror("open loop");
		exit(1);
	}

	if ((backfd = open(backfile, O_RDONLY)) < 0) {
		perror("open backfile");
		exit(1);
	}

	ret = ioctl(fd, LOOP_CHANGE_FD, backfd);
	if (ret) {
		perror("change fd failed");
		exit(1);
	}

	memset(&lo_info, 0, sizeof(lo_info));
	strncpy((char*)lo_info.lo_file_name, backfile, LO_NAME_SIZE);
	lo_info.lo_file_name[LO_NAME_SIZE - 1] = 0;
	
	ret = ioctl(fd, LOOP_SET_STATUS64, &lo_info);
	if (ret) {
		perror("set status failed");
		exit(1);
	}
	
	close(backfd);
	close(fd);
	printf("change fd: success\n");
}

void show_rblocked()
{
	int ret, pos, fd;
	char buf[4096];

	if ((fd = open(device, O_RDONLY)) < 0) {
		perror("open loop");
		exit(1);
	}

	pos = 0;
	ret = read(fd, buf, 4096);
	while (ret >= 0) {
		pos += ret;
		if (ret == 0) {
			if (pos == 0)
				break;
			else
				lseek(fd, 0, SEEK_SET);
		}
		write(STDOUT_FILENO, buf, ret);
		ret = read(fd, buf, 4096);
	}
	close(fd);
}

int main(int argc, char **argv)
{
	int cmd = CMD_INFO;	
	char *p, c;
	progname = argv[0];

	if ((p = strrchr(progname, '/')) != NULL)
		progname = p+1;

	while ((c = getopt(argc, argv, "w:x:m:r:s:b:t:c:d:v")) != -1) {
		switch (c) {
		case 'm':
			cmd = CMD_MAP_INFO;
			mapfile = optarg;
			break;
		case 'w':
			cmd = CMD_WLIMIT_START;
			device = optarg;
			break;
		case 'x':
			cmd = CMD_WLIMIT_STOP;
			device = optarg;
			break;
		case 'r':
			cmd = CMD_RLIMIT_START;
			device = optarg;
			break;
		case 's':
			cmd = CMD_RLIMIT_STOP;
			device = optarg;
			break;
		case 't':
			cmd = CMD_SYNC_BLOCK;
			device = optarg;
			break;
		case 'c':
			cmd = CMD_CHANGE;
			device = optarg;
			break;
		case 'd':
			cmd = CMD_SHOW_RBLOCKED;
			device = optarg;
			break;
		case 'b':
			if (1 != sscanf(optarg, "%d", &block_index))
				block_index = -1;
			else if (block_index < 0)
				block_index = -1;
			break;
		case 'v':
			verbose = 1;
			break;
		default:
			usage();
		}
	}

	if (argc == 1)
		usage();

	if (cmd == CMD_WLIMIT_STOP || cmd == CMD_RLIMIT_START) {
		if (argc != optind + 1)
			usage();
		mapfile = argv[optind];
	}
	else if (cmd == CMD_SYNC_BLOCK) {
		if (argc != optind + 1 || block_index < 0)
			usage();
		backfile = argv[optind];
	}
	else if (cmd == CMD_CHANGE) {
		if (argc != optind + 1)
			usage();
		backfile = argv[optind];
	}


	if (cmd == CMD_INFO)
		show_info();
	else if (cmd == CMD_WLIMIT_START)
		wlimit_start();
	else if (cmd == CMD_WLIMIT_STOP)
		wlimit_stop();
	else if (cmd == CMD_RLIMIT_START)
		rlimit_start();
	else if (cmd == CMD_RLIMIT_STOP)
		rlimit_stop();
	else if (cmd == CMD_MAP_INFO)
		show_map_info();
	else if (cmd == CMD_SYNC_BLOCK)
		sync_block();
	else if (cmd == CMD_CHANGE)
		change_fd();
	else if (cmd == CMD_SHOW_RBLOCKED)
		show_rblocked();

	return 0;
}
